#!/usr/bin/env python3
# coding=utf-8

import rospy
from sensor_msgs.msg import JointState
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# 全局数据缓冲
time_vals = []
positions_vals = [[] for i in range(7)]

def joint_callback(msg):
    """ROS Callback，获取每一个关节的数据。"""
    if not time_vals:
        time_vals.append(0.0)
    else:
        time_vals.append(time_vals[-1] + 0.05)  # 假如数据是20 Hz发布
    
    for i in range(7):
        positions_vals[i].append(msg.position[i])

def animate(i):
    """每隔一定时间被 FuncAnimation调用，重画数据。"""
    if not time_vals:
        return
    
    for idx, ax in enumerate(axs.flatten()):

        ax.cla()
        
        if idx < 7:
            ax.plot(time_vals, positions_vals[idx], label=f'Joint {idx+1}')
            ax.set_xlabel('Time (sec)')
            ax.set_ylabel('Position (radian)')
            ax.legend()
            ax.grid()
        else:
            ax.axis('off')  # 最后一个空着

def main():
    """ROS Main."""
    rospy.init_node('joints_plotter', anonymous=True)
    rospy.Subscriber('/joint_states', JointState, joint_callback)

    # 设置 matplotlib
    global fig, axs
    fig, axs = plt.subplots(4, 2, figsize=(10, 12))
    ani = animation.FuncAnimation(fig, animate, interval=50)  # 每50ms刷新
    
    plt.tight_layout()
    plt.show()
    rospy.spin()


if __name__ == '__main__':
    main()
